# -*- encoding: utf-8 -*-
"""
根据所提供的的数据，生成wav.scp、segments、utt2spk、text
"""
import os,sys
import re
import string
import wave
import argparse


def get_args():
    parser = argparse.ArgumentParser(description="""data preparation long style dataset""")
    parser.add_argument('--wavdir', type=str, required=True, default='',
                        help='input wave dataset dir')
    parser.add_argument('--outdir', type=str, required=True, default='', 
                        help='output data dir')
    parser.add_argument('--is-remove-tags', type=bool, default=True, 
                        help="remove tags, for examples:[ENS] or [噪声]")
    args = parser.parse_args()
    return args


def remove_punctuation(text):
    """移除标点符号"""
    PUNCTIUATION = u"[！？。＂＃＄％＆＇（）＊＋－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏" + string.punctuation + ']+'
    rule = re.compile(PUNCTIUATION)
    text = re.sub(rule, "", text)
    return text


def remove_tags(text):
    """移除特殊标识符
    如：[ENS]、[LAUGH] ......
    """
    rule = re.compile('(\[[^\]|^\[]*\])')
    text = re.sub(rule, "", text)
    return text
    

def find_by_pattern(path, pattern='.wav'):
    result_dict = {}
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith(pattern):
                result_dict[f] = os.path.abspath(os.path.join(root, f))
    return result_dict


def preprocess_data(wav_dir, txt_dir, output_dir, is_remove_tags=True):
    """
    根据数据集的数据结构（文件结构），生成wav.scp、segments、utt2spk、text
    
    dataset
    ├── README.txt
    ├── TXT
    │   ├── CTS-CN-F2F-2017-11-14_190.txt
    │   ├── CTS-CN-F2F-2017-11-14_215.txt
    |   .
    │   └── CTS-CN-F2F-2019-11-15-185.txt
    └── WAV
        ├── CTS-CN-F2F-2017-11-14_190.wav
        ├── CTS-CN-F2F-2017-11-14_215.wav
        .
        └── CTS-CN-F2F-2019-11-15-185.wav
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    wavs_dict = find_by_pattern(wav_dir, pattern='.wav')
    txts_dict = find_by_pattern(txt_dir, pattern='.txt')
    
    wav_fout = open(os.path.join(output_dir, 'wav.scp') ,'w')
    utt_fout = open(os.path.join(output_dir, 'utt2spk') ,'w')
    seg_fout = open(os.path.join(output_dir, 'segments') ,'w')
    text_fout = open(os.path.join(output_dir, 'text') ,'w')
    
    for txt_item in txts_dict:
        file_name = txt_item.replace('.txt', '.wav')
        if file_name in wavs_dict:
            wav_fout.write('%s %s\n'%(file_name, wavs_dict[file_name]))
            with open(txts_dict[txt_item], 'r') as fin:
                for line in fin:
                    line_split = line.strip().split('\t')
                    assert len(line_split) == 4, "%s"%line_split
                    beg, end = eval(line_split[0])
                    if beg < end:
                        seg_text = line_split[3]
                        
                        
                        # 移除[]标识符和标点
                        if is_remove_tags:
                            seg_text = remove_tags(seg_text)
                        seg_text = remove_punctuation(seg_text)
                        if len(seg_text):
                            uttid = '%s-%s-%s' % (file_name, '%07d' % round(beg*1000), '%07d' % round(end*1000))
                            utt_fout.write('%s %s\n'%(uttid, uttid))
                            seg_fout.write('%s %s %s %s\n'%(uttid, file_name, str(beg), str(end)))
                            text_fout.write('%s %s\n'%(uttid, seg_text))
                        else:
                            print("Info: {%s} in %s"%(line.strip(), txt_item))
                    else:
                        print("Error: {%s} in %s"%(line.strip(), txt_item))

    wav_fout.close()
    utt_fout.close()
    seg_fout.close()
    text_fout.close()
    print("Prepare done")


def main():
    args = get_args()
    print(args.__dict__)
    
    wav_dir = os.path.join(args.wavdir, 'WAV') 
    txt_dir = os.path.join(args.wavdir, 'TXT')
    if not os.path.exists(wav_dir) or not os.path.exists(txt_dir):
        raise IOError("%s not found WAV or TXT"%(args.wavdir))
    output_dir = args.outdir
    preprocess_data(wav_dir, txt_dir, output_dir, is_remove_tags=args.is_remove_tags)
    
    
if __name__ == '__main__':
    main()
